package com.gearbox.platform.slurm.service;

import com.gearbox.core.configuration.SystemConfiguration;
import com.gearbox.core.constant.NodeStatusEnum;
import com.gearbox.core.constant.SlurmInstanceStatusEnum;
import com.gearbox.core.model.Node;
import com.gearbox.core.service.NodeService;
import com.gearbox.platform.slurm.util.SlurmCmdExecuteResultParser;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Component;

import java.util.Collection;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.stream.Collectors;

@Component
@ConditionalOnProperty(name = "system.type", havingValue = "slurm")
public class SlurmNodeService implements NodeService {
    @Autowired
    SystemConfiguration systemConfiguration;

    private static final String SCALE_IN_REASON = "scale in node";

    @Override
    public List<Node> listUnstableNodes() {
        return getNodesByPartition(systemConfiguration.getVariablePartition()).stream()
            .filter(node ->
                !systemConfiguration.getStableNodes().contains(node.getName()))
            .collect(Collectors.toList());
    }

    @Override
    public List<Node> listUnstableDrainNodes() {
        return getNodesByPartition(systemConfiguration.getVariablePartition()).stream()
            .filter(node ->
                node.getStatus().equals(NodeStatusEnum.DRAIN)
                    && !systemConfiguration.getStableNodes().contains(node.getName()))
            .collect(Collectors.toList());
    }

    @Override
    public List<Node> listUnstableIdleNodes() {
        return getNodesByPartition(systemConfiguration.getVariablePartition()).stream()
            .filter(node ->
                node.getStatus().equals(NodeStatusEnum.IDLE)
                    && !systemConfiguration.getStableNodes().contains(node.getName()))
            .collect(Collectors.toList());
    }

    Set<Node> getNodesByPartition(String partition) {
        List<String> nodesInfo = SlurmCommandExecutor.getNodesInfoByPartition(partition);
        return nodesInfo.stream()
            .map(SlurmCmdExecuteResultParser::parseSInfoLine)
            .flatMap(Collection::stream)
            .collect(Collectors.toSet());
    }

    @Override
    public void remove(Node node) {
        deleteSlurmNode(node);
        deleteHostRecord(node);
    }

    void deleteSlurmNode(Node node) {
        SlurmCommandExecutor.deleteNodeByName(node.getName());
    }

    public void deleteHostRecord(Node node) {
        SlurmCommandExecutor.deleteHostsRecordByKeyword(node.getName());
    }

    @Override
    public void setToDrain(String nodeName) {
        SlurmCommandExecutor.updateNodeStatus(
            nodeName, SlurmInstanceStatusEnum.DRAIN.name().toUpperCase(Locale.ROOT), SCALE_IN_REASON);
    }

    @Override
    public int getUserNodeQuotas(String userName) {
        List<String> quotas = SlurmCommandExecutor.getQuotasByUser(userName);
        return quotas.stream().mapToInt(Integer::parseInt).sum();
    }

    @Override
    public boolean isWorking(String nodeName) {
        List<String> jobInfos = SlurmCommandExecutor.getJobInfosByNodeName(nodeName);
        for (String line : jobInfos) {
            if (line.contains(nodeName)) {
                return true;
            }
        }
        return false;
    }
}
